#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import json
import bbki
import elemlib
import gstage4
import platform
import strict_hdds
import strict_hwcfg
import strict_portage
from fm_util import Util
from fm_param import FmConst
from helper_bbki import BbkiWrapper
from helper_bbki import BbkiCfgFileUtil
from helper_dyncfg import DynCfgModifier
from helper_pkg_warehouse import PkgWarehouse
from helper_remote import RemoteOrLocalBuild
from stage4_gentoo import FpemudOs


class FmSysUpdater:

    def __init__(self, param):
        self.param = param
        self.infoPrinter = self.param.infoPrinter

    def update(self, bSync):
        self.infoPrinter.printInfo(">> Preparing...")
        machine = None
        layout = None
        bootEntrySpec = None
        if True:
            self.param.sysChecker.init(False, False)
            self.param.sysChecker.basicCheck()

            DynCfgModifier.update()

            machine = strict_hwcfg.probe()

            if self.param.runMode in ["normal", "setup"]:
                layout = strict_hdds.get_storage_layout()

            bbkiObj = BbkiWrapper(machine, layout=layout)

            bootEntrySpec = None
            if not bSync:
                if self.param.runMode in ["normal", "setup"]:
                    be = bbkiObj.getPendingBootEntry()
                    if be is not None:
                        bootEntrySpec = be.spec
            if bootEntrySpec is None:
                bootEntrySpec = self._getBbkiSpec(machine)

            # check python & ruby version
            pkgwh = PkgWarehouse()
            pythonVer = pkgwh.pythonGetCurrentVersion()
            newestPythonVer = pkgwh.pythonGetValidNewestVersion(pythonVer) if pythonVer is not None else None
            rubyVer = pkgwh.rubyGetCurrentVersion()
            newestRubyVer = pkgwh.rubyGetValidNewestVersion(rubyVer) if rubyVer is not None else None

            # update some files in /etc/portage directory
            self._updateSomeFiles(machine)

            # set system to unstable status
            if self.param.runMode in ["normal", "setup"]:
                if bbkiObj.isStable():
                    with strict_hdds.BootDirWriter(layout):
                        bbkiObj.setStable(False)

            print("")

        with RemoteOrLocalBuild(self.infoPrinter) as obj:
            # sync
            with obj.getExecContext(syncDownType="system") as ctx:
                ctx.exec(os.path.join(FmConst.libexecDir, "op-sync.py"),
                         str(self.infoPrinter.getIndent()),
                         Util.boolToStr01(bSync),
                         json.dumps(bootEntrySpec.to_dict()))

            # fetch and build
            if True:
                kernelVersion = None
                with bbkiObj.getKernelInstaller(bootEntrySpec) as kernelInstaller:
                    kernelVersion = kernelInstaller.get_kernel_version()

                # install kernel, initramfs and bootloader
                with strict_hdds.BootDirWriter(layout):
                    self.infoPrinter.printInfo(">> Installing Linux-%s-%s..." % (bootEntrySpec.arch.value, kernelVersion))
                    with obj.getExecContext(syncDownType="kernel") as ctx:
                        bKernelBuildNeeded = True
                        if not bSync and self.param.runMode in ["normal", "setup"]:
                            oldBootEntry = bbkiObj.getPendingBootEntry()
                            if oldBootEntry is not None and oldBootEntry.has_kernel_files() and oldBootEntry.spec != bootEntrySpec:
                                bKernelBuildNeeded = False

                        resultFile = os.path.join(self.param.tmpDir, "result.txt")
                        ctx.exec(os.path.join(FmConst.libexecDir, "op-install-kernel.py"),
                                 str(self.infoPrinter.getIndent()),
                                 json.dumps(bootEntrySpec.to_dict()),
                                 Util.boolToStr01(bKernelBuildNeeded),
                                 oldBootEntry.postfix if not bKernelBuildNeeded else "",
                                 oldBootEntry.kernel_config_filepath if not bKernelBuildNeeded else "",
                                 resultFile)
                        print("")

                    self.infoPrinter.printInfo(">> Creating initramfs...")
                    if True:
                        if self.param.runMode in ["normal", "setup"]:
                            bbkiObj.updateInitramfs()
                        else:
                            print("WARNING: Running in \"%s\" mode, do NOT create initramfs!!!" % (self.param.runMode))
                        print("")

                    self.infoPrinter.printInfo(">> Updating boot-loader...")
                    if self.param.runMode in ["normal", "setup"]:
                        bbkiObj.updateBootloader()
                    else:
                        print("WARNING: Running in \"%s\" mode, do NOT maniplate boot-loader!!!" % (self.param.runMode))
                    print("")

                # synchronize boot partitions
                if layout is not None and layout.name in ["efi-btrfs", "efi-bcache-btrfs", "efi-bcachefs"]:
                    dstList = layout.get_pending_esp_list()
                    if len(dstList) > 0:
                        with self.infoPrinter.printInfoAndIndent(">> Synchronizing boot partitions..."):
                            for dst in dstList:
                                self.infoPrinter.printInfo("        - %s to %s..." % (layout.get_esp(), dst))
                                layout.sync_esp(dst)
                        print("")

            with obj.getExecContext(syncDownType="system") as ctx:
                # emerge @world
                self.infoPrinter.printInfo(">> Updating @world...")
                ctx.exec(os.path.join(FmConst.libexecDir, "op-emerge-world.py"))

                # update versions if neccessary
                if pythonVer is not None:
                    self.infoPrinter.printInfo(">> Selecting newest versions...")
                    ctx.exec(os.path.join(FmConst.libexecDir, "op-update-versions.py"),
                             str(self.infoPrinter.getIndent()),
                             pythonVer if pythonVer is not None else "",
                             newestPythonVer if newestPythonVer is not None else "",
                             rubyVer if rubyVer is not None else "",
                             newestRubyVer if newestRubyVer is not None else "")

                # re-emerge all "-9999" packages
                self.infoPrinter.printInfo(">> Updating all \"-9999\" packages...")
                ctx.exec(os.path.join(FmConst.libexecDir, "op-emerge-9999.py"))

        # FIXME
        # with self.infoPrinter.printInfoAndIndent(">> Updating elements..."):
        #     with elemlib.Library(system_library_or_user_library=True) as sysLib:
        #         with self.infoPrinter.printInfoAndIndent("- Synchronize metadata..."):
        #             for etype in FmConst.supportElementTypes:
        #                 self.infoPrinter.printInfo("- Type %s..." % (etype))
        #                 sysLib.sync(etype)
        #         with self.infoPrinter.printInfoAndIndent("- Update cached elements..."):
        #             for etype in FmConst.supportElementTypes:
        #                 for elem in sysLib.get_elements(etype, "*", "*"):
        #                     self.infoPrinter.printInfo("- Type: %s, Source: %s, Index: %s..." % (elem.etype, elem.esource, elem.eindex))
        #                     cachedPartNameList = [x.pname for x in elem.get_parts() if x.get_local_path() is not None]
        #                     elem.cache(cachedPartNameList)

    def installPackage(self, pkgName, bTest):
        self.infoPrinter.printInfo(">> Preparing...")
        if True:
            self.param.sysChecker.init(False, False)
            self.param.sysChecker.basicCheck(repositoryOverlayCheck=True)

            DynCfgModifier.update()

            machine = strict_hwcfg.probe()

            # check python & ruby version
            pkgwh = PkgWarehouse()
            pythonVer = pkgwh.pythonGetCurrentVersion()
            rubyVer = pkgwh.rubyGetCurrentVersion()

            # update some files in /etc/portage directory
            self._updateSomeFiles(machine)

            print("")

        with RemoteOrLocalBuild(self.infoPrinter) as obj:
            # add pre-enabled overlays and packages
            with obj.getExecContext(syncDownType="system") as ctx:
                ctx.exec(os.path.join(FmConst.libexecDir, "op-sync.py"),
                         str(self.infoPrinter.getIndent()),
                         Util.boolToStr01(False),               # specify False here so that no synchronization would happen
                         "")

            # emerge package
            with obj.getExecContext(syncDownType="system") as ctx:
                self.infoPrinter.printInfo(">> Installing %s..." % (pkgName))
                ctx.exec(os.path.join(FmConst.libexecDir, "op-emerge-package.py"),
                         pkgName,
                         Util.boolToStr01(bTest))

        # FIXME: will version be auto raised?
        assert pythonVer is None or PkgWarehouse().pythonGetCurrentVersion() == pythonVer
        assert rubyVer is None or PkgWarehouse().rubyGetCurrentVersion() == rubyVer

    def stablize(self):
        machine = strict_hwcfg.probe()
        layout = strict_hdds.get_storage_layout()
        bbkiObj = BbkiWrapper(machine, layout=layout)

        self.infoPrinter.printInfo(">> Stablizing...")
        with strict_hdds.BootDirWriter(layout):
            bbkiObj.setStable(True)
        print("")

        if layout.name in ["efi-btrfs", "efi-bcache-btrfs", "efi-bcachefs"]:
            dstList = layout.get_pending_esp_list()
            if len(dstList) > 0:
                with self.infoPrinter.printInfoAndIndent(">> Synchronizing boot partitions..."):
                    for dst in dstList:
                        self.infoPrinter.printInfo("        - %s to %s..." % (layout.get_esp(), dst))
                        layout.sync_esp(dst)
                print("")

    def updateAfterHddAddOrRemove(self, layout, bbkiObj):
        pendingBe = bbkiObj.getPendingBootEntry()
        if pendingBe is None:
            raise Exception("No boot entry found!")

        # re-create initramfs
        with strict_hdds.BootDirWriter(layout):
            self.infoPrinter.printInfo(">> Recreating initramfs...")
            bbkiObj.updateInitramfs()
            print("")

            self.infoPrinter.printInfo(">> Updating boot-loader...")
            bbkiObj.updateBootloader()
            print("")

        # synchronize boot partitions
        if layout.name in ["efi-btrfs", "efi-bcache-btrfs", "efi-bcachefs"]:
            dstList = layout.get_pending_esp_list()
            if len(dstList) > 0:
                with self.infoPrinter.printInfoAndIndent(">> Synchronizing boot partitions..."):
                    for dst in dstList:
                        self.infoPrinter.printInfo("        - %s to %s..." % (layout.get_esp(), dst))
                        layout.sync_esp(dst)
                print("")

    def _parseKernelBuildResult(self, result):
        lines = result.split("\n")
        lines = [x.rstrip() for x in lines if x.rstrip() != ""]
        assert len(lines) == 2
        return (lines[0] != "0", lines[1])       # (kernelBuilt, postfix)

    def _getBbkiSpec(self, machine):
        tKernelType = "linux"
        tKernelParams = {}
        if os.path.exists(FmConst.bbkiKernelFile):
            tKernelType, tKernelParams = BbkiCfgFileUtil.parseKernelCfgFile(FmConst.bbkiKernelFile)
        else:
            pass                                # use default

        tKernelAddonDict = {}
        if os.path.exists(FmConst.bbkiKernelAddonDir):
            for fn in os.listdir(FmConst.bbkiKernelAddonDir):
                fullfn = os.path.join(FmConst.bbkiKernelAddonDir, fn)
                for addonName, addonParams, bAdd in BbkiCfgFileUtil.parseKernelAddonCfgFile(fullfn):
                    if bAdd:
                        tKernelAddonDict[addonName] = addonParams
                    else:
                        tKernelAddonDict.pop(addonName)
        else:
            tKernelAddonDict.update({           # use default
                "pnp": {
                    "peripheral": None,
                },
                "noatime_default": {},
                "wireless_regdb": {},
                "linux_firmware": {
                    "strict": None,
                },
                "fedora_alsa_firmware": {
                    "strict": None,
                },
                "riptide_firmware": {
                    "strict": None,
                },
                "ipw2100_firmware": {
                    "strict": None,
                },
                "ipw2200_firmware": {
                    "strict": None,
                },
                "storage_layouts": {
                    # "efi-bcachefs": None,
                    "efi-bcache-btrfs": None,
                    "efi-btrfs": None,
                    "efi-ext4": None,
                    "bios-ext4": None,
                },
            })

            if self.param.runMode == "normal":
                cpuVendor, _ = Util.hwcfgGetCpuVendorModel(machine)
                if cpuVendor == "intel":
                    tKernelAddonDict.update({
                        "x86_32": {},
                        "x86_intel": {},
                    })
                elif cpuVendor == "amd":
                    tKernelAddonDict.update({
                        "x86_32": {},
                        "x86_amd": {},
                    })
                    # FIXME: it seems zenpower3 git website is not accessable anymore
                    # if "ryzen" in cpu.get_model():
                    #     standardBuf += "zenpower3\n"
                else:
                    assert False
            elif self.param.runMode in ["setup", "prepare"]:
                tKernelAddonDict.update({
                    "x86_32": {},
                    "x86_intel": {},
                    "x86_amd": {},
                })
            else:
                assert False

            if self.param.runMode == "normal":
                mobo = Util.hwcfgGetMoboComponent(machine)
                if mobo is not None and mobo.get_numa():
                    tKernelAddonDict.update({
                        "x86_numa": {},
                    })
            elif self.param.runMode in ["setup", "prepare"]:
                pass
            else:
                assert False

            # bluetooth
            #     standardBuf += "rtl8761\n"

            tKernelAddonDict.update({
                "vm_as_guest": {},
                # "virtualbox_as_guest": {},
            })

            tKernelAddonDict.update({
                "iwd": {},
                "ddcci_driver": {},
                "vhba_module": {},
                "vm_as_host": {},
                # "virtualbox_as_host": {},
                "wine": {},
                "display_full": {},
                "desktop_functions": {},
                "xpadneo": {
                    "rumble_attenuation": "100",      # disable rumble
                    "disable_shift_mode": "1",        # don't give xbox button extra functions
                },
            })

        return bbki.BootEntrySpec({
            "arch": platform.machine(),
            "kernel_type": tKernelType,
            "kernel_params": tKernelParams,
            "kernel_addon_params": tKernelAddonDict,
            "initramfs_params": {},
        })

    def _updateSomeFiles(self, machine):
        portageConfigDir = strict_portage.cfg.PortageConfigDir()
        fpemudOs = FpemudOs()
        pkgwh = PkgWarehouse()

        ts = gstage4.TargetSettings()
        fpemudOs.update_target_settings(ts)

        def __errCb(checkCode, message):
            raise Exception(message)

        # /etc/portage/repo.postsync.d directory
        if len(ts.repo_postsync_patch_directories) > 0:
            reposPostSyncDir = portageConfigDir.get_repo_postsync_dir_obj()
            c = reposPostSyncDir.create_checker(auto_fix=True, error_callback=__errCb)
            fpemudOs.checkPortageConfigDirByTargetSettings(reposPostSyncDir, c, ts, True, __errCb)

        # /etc/portage/package.license file
        if len(ts.pkg_license) > 0:
            packageLicense = portageConfigDir.get_package_license_obj(file_or_dir=True)
            with packageLicense.create_checker(auto_fix=True, error_callback=__errCb) as c:
                fpemudOs.checkPortageConfigDirByTargetSettings(packageLicense, c, ts, True, __errCb)

        # /etc/portage/package.mask directory
        if True:
            packageMask = portageConfigDir.get_package_mask_obj(file_or_dir=False)
            c = packageMask.create_checker(auto_fix=True, error_callback=__errCb)
            fpemudOs.checkPortageConfigDirByTargetSettings(packageMask, c, ts, True, __errCb)

        # /etc/portage/package.use directory
        if True:
            packageUse = portageConfigDir.get_package_use_obj(file_or_dir=False)
            c = packageUse.create_checker(auto_fix=True, error_callback=__errCb)
            fpemudOs.checkPortageConfigDirByTargetSettings(packageUse, c, ts, True, __errCb)

            # /etc/portage/package.use/30-hardware
            pkgwh.refreshHardwareUseFlags(machine)

            # /etc/portage/package.use/90-python-targets
            pkgwh.pythonRefreshTargetUseFlags()

            # /etc/portage/package.use/90-ruby-targets
            pkgwh.rubyRefreshTargetUseFlags()

        # /etc/portage/package.env and /etc/portage/env directory
        if True:
            packageEnv = portageConfigDir.get_package_env_obj(file_or_dir=False)
            with packageEnv.create_checker(auto_fix=True, error_callback=__errCb) as c:
                c.check_self()
                fpemudOs.checkPortageConfigDirByTargetSettings(packageEnv, c, ts, True, __errCb)
